import happi

import numpy as np
import matplotlib.pyplot as plt
plt.ion()
ln = np.log

t0 = 1e-12 # seconds


Te=np.array([ 1.00000000e-03, 1.26638017e-03, 1.60371874e-03, 2.03091762e-03, 2.57191381e-03,
	3.25702066e-03, 4.12462638e-03, 5.22334507e-03, 6.61474064e-03, 8.37677640e-03,
	1.06081836e-02, 1.34339933e-02, 1.70125428e-02, 2.15443469e-02, 2.72833338e-02,
	3.45510729e-02, 4.37547938e-02, 5.54102033e-02, 7.01703829e-02, 8.88623816e-02,
	1.12533558e-01, 1.42510267e-01, 1.80472177e-01, 2.28546386e-01, 2.89426612e-01,
	3.66524124e-01, 4.64158883e-01, 5.87801607e-01, 7.44380301e-01, 9.42668455e-01,
	1.19377664e+00, 1.51177507e+00, 1.91448198e+00, 2.42446202e+00, 3.07029063e+00,
	3.88815518e+00, 4.92388263e+00, 6.23550734e+00, 7.89652287e+00, 1.00000000e+01])
Z_ = {
"H":np.array([ 7.88233875e-04, 3.92730202e-03, 1.44293662e-02, 4.14461367e-02, 9.71150139e-02,
	1.91342157e-01, 3.24200383e-01, 4.82830926e-01, 6.38408095e-01, 7.48811369e-01,
	8.09442340e-01, 8.46076403e-01, 8.74775837e-01, 8.99450346e-01, 9.20308767e-01,
	9.37378760e-01, 9.51022865e-01, 9.61776860e-01, 9.70189775e-01, 9.76748122e-01,
	9.81854289e-01, 9.85829673e-01, 9.88926462e-01, 9.91340840e-01, 9.93224877e-01,
	9.94696334e-01, 9.95846456e-01, 9.96746026e-01, 9.97450031e-01, 9.98001251e-01,
	9.98433016e-01, 9.98771323e-01, 9.99036471e-01, 9.99244325e-01, 9.99407293e-01,
	9.99535086e-01, 9.99635307e-01, 9.99713912e-01, 9.99775568e-01, 9.99823931e-01]),
"Al":np.array([ 0.04609447, 0.11788394, 0.25767174, 0.47789452, 0.75850024, 1.05706387, 1.34063202,
	1.60385045, 1.86151729, 2.13395035, 2.43931737, 2.79212759, 3.20334245, 3.68008809,
	4.22496502, 4.83564283, 5.50532988, 6.22401652, 6.97957234, 7.75759139, 8.53970481,
	9.30169869, 10.0143301 , 10.64902124, 11.18621051, 11.62046189, 11.95911407, 12.21684622,
	12.41025716, 12.55452124, 12.66206557, 12.74246417, 12.80285101, 12.8484551 , 12.88308641,
	12.90952262, 12.92979784, 12.9454115 , 12.95747717, 12.9668282 ]),
"Zn":np.array([ 0.05025318, 0.14303339, 0.32539979, 0.61258922, 0.98342111, 1.38727693, 1.77528062,
	2.12762036, 2.45521265, 2.78382598, 3.14028677, 3.54789603, 4.02720893, 4.5978342 ,
	5.27948136, 6.09177991, 7.0526969 , 8.17557912, 9.46538605, 10.91548677, 12.50678131,
	14.20995663, 15.98953018, 17.8067511 , 19.61919475, 21.37783514, 23.02586986, 24.50487373,
	25.76902954, 26.79891154, 27.6043477 , 28.21561975, 28.67090493, 29.00688683, 29.25427869,
	29.43688072, 29.57235323, 29.67350816, 29.74955275, 29.80709617]),
"Au":np.array([ 0.42095719, 0.70478429, 1.06769482, 1.48588245, 1.9242277 , 2.34976608, 2.74561441,
	3.11604664, 3.4807232 , 3.86482211, 4.2924045 , 4.78473649, 5.36143821, 6.04235616,
	6.84929956, 7.80752012, 8.9469866 , 10.30340694, 11.91873264, 13.84057796, 16.11965772,
	18.80418339, 21.9305915 , 25.51153958, 29.52473287, 33.90815976, 38.56573939, 43.38130946,
	48.23275232, 52.99823957, 57.55375982, 61.76990959, 65.52021161, 68.70651127, 71.28856372,
	73.29422978, 74.80222575, 75.91222479, 76.72040917, 77.30714207]),
}

colors = ["b","k","r","g"]

fig = plt.figure(3, figsize=(6,3.5))

for element in ["H", "Al", "Zn", "Au"]:
	print("Analyzing "+element)
	
	S=happi.Open("ionization_equilibrium"+element)
	
	npoints = S.namelist.npoints
	every = S.namelist.DiagParticleBinning[0].every
	ts = int(t0 * S.namelist.Main.reference_angular_frequency_SI/S.namelist.Main.timestep/every) # timestep at 1ps
    
	Z = []
	Zfinal = []
	T = []
	Tfinal = []
	for i in range(npoints):
		D = S.ParticleBinning("#"+str(4*i+2)+"/#"+str(4*i+3),sum={"x":"all"}, units=["ps"],marker=".")
		Z.append( D )
		Zfinal.append( np.array(D.getData())[ts] )
		
		D = S.ParticleBinning("#"+str(4*i+0)+"/#"+str(4*i+1),sum={"x":"all"}, units=["ps"],marker=".")
		T.append( D )
		Tfinal.append( np.array(D.getData())[ts] *(511.*2./3.) )
		
		
	
	#happi.multiPlot(*Z, figure=1, skipAnimation=True)
	
	#happi.multiPlot(*T, figure=2, skipAnimation=True)
	
	color = colors.pop()
	
	plt.figure(3)
	plt.loglog(Tfinal, Zfinal, "o"+color, markeredgecolor="none")
	plt.plot(Te, Z_[element], "-"+color, label=element)

ax = plt.gca()
ax.set_position([0.13,0.17,0.85,0.77])
fig.set_facecolor("w")
ax.add_artist(ax.legend(loc=2, prop={'size':11}))
ax.xaxis.labelpad = 0
ax.yaxis.labelpad = 0
ax.set_xlabel("Final temperature (keV)")
ax.set_ylabel("Average ion charge at $t=1$ ps")
ax.set_xlim(3e-3, 5)
ax.set_ylim(0.1, 100)

l1, = plt.plot([0],"ok", label="Smilei")
l2, = plt.plot([0],"-k", label="Thomas-Fermi")
plt.legend(handles=[l1, l2], loc=4, prop={'size':11})


